import json
import csv
import os
import re
from pathlib import Path

def merge_json_files_to_csv():
    # 定义目录路径
    directory = Path('./zero_gen')
    
    # 创建CSV文件
    csv_file_path = directory / 'merged_data.csv'
    
    # 收集所有需要处理的JSON文件
    json_files = list(directory.glob('generated_normal_*.json'))
    
    # 准备CSV数据
    csv_data = []
    
    # 处理每个JSON文件
    for json_file in json_files:
        # 从文件名提取类别（去掉前缀和扩展名）
        category = json_file.stem.replace('generated_normal_', '')
        
        # 读取JSON文件
        with open(json_file, 'r', encoding='utf-8') as f:
            data = json.load(f)
            
        # 处理每个条目
        for item in data:
            if 'answer' in item:
                # 添加到CSV数据中
                csv_data.append({
                    'prompt': item.get('prompt', ''),  # 如果没有prompt字段，则为空字符串
                    'answer': item['answer'],
                    'category': category
                })
    
    # 写入CSV文件
    with open(csv_file_path, 'w', newline='', encoding='utf-8') as csvfile:
        fieldnames = ['prompt', 'answer', 'category']
        writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
        
        # 写入表头
        writer.writeheader()
        
        # 写入数据
        for row in csv_data:
            writer.writerow(row)
    
    print(f"成功合并 {len(json_files)} 个JSON文件到 {csv_file_path}")
    print(f"总共处理了 {len(csv_data)} 条记录")

if __name__ == "__main__":
    merge_json_files_to_csv()